"""Token使用统计（主模型调用计费）"""
"""Token计数器（统计请求/响应token数量）"""
import time
from typing import Dict, Any
from fin_senti_entity_platform.utils.logger import Logger
from fin_senti_entity_platform.utils.constants import PERFORMANCE_METRICS


class TokenCounter:
    """Token计数器类，用于统计和监控token使用情况"""
    
    def __init__(self):
        """初始化Token计数器"""
        self.logger = Logger().get_logger(__name__)
        self.total_input_tokens = 0
        self.total_output_tokens = 0
        self.request_count = 0
        self.last_reset_time = time.time()
        self._lock = self._get_threading_lock()
        
    def _get_threading_lock(self):
        """获取线程锁，避免循环导入"""
        import threading
        return threading.Lock()
    
    def update_count(self, input_tokens: int, output_tokens: int):
        """更新token计数"""
        with self._lock:
            self.total_input_tokens += input_tokens
            self.total_output_tokens += output_tokens
            self.request_count += 1
    
    def get_stats(self) -> Dict[str, Any]:
        """获取token统计信息"""
        with self._lock:
            elapsed_time = time.time() - self.last_reset_time
            if elapsed_time == 0:
                elapsed_time = 1  # 避免除零错误
            
            return {
                'total_input_tokens': self.total_input_tokens,
                'total_output_tokens': self.total_output_tokens,
                'total_tokens': self.total_input_tokens + self.total_output_tokens,
                'request_count': self.request_count,
                'avg_input_tokens_per_request': self.total_input_tokens / self.request_count if self.request_count > 0 else 0,
                'avg_output_tokens_per_request': self.total_output_tokens / self.request_count if self.request_count > 0 else 0,
                'tokens_per_second': (self.total_input_tokens + self.total_output_tokens) / elapsed_time,
                'elapsed_time': elapsed_time
            }
    
    def reset(self):
        """重置计数器"""
        with self._lock:
            self.total_input_tokens = 0
            self.total_output_tokens = 0
            self.request_count = 0
            self.last_reset_time = time.time()
            self.logger.info("Token计数器已重置")
    
    def estimate_cost(self, input_price_per_1m: float = 0.01, output_price_per_1m: float = 0.03) -> float:
        """估算使用成本
        
        Args:
            input_price_per_1m: 每百万输入token的价格
            output_price_per_1m: 每百万输出token的价格
            
        Returns:
            float: 估算成本
        """
        with self._lock:
            input_cost = (self.total_input_tokens / 1_000_000) * input_price_per_1m
            output_cost = (self.total_output_tokens / 1_000_000) * output_price_per_1m
            return input_cost + output_cost


# 创建全局实例
global_token_counter = TokenCounter()


def get_token_counter() -> TokenCounter:
    """获取Token计数器实例"""
    return global_token_counter